#pragma region Copyright
/// <copyright>
/// Copyright (c) 2007 Ilia Shramko.
/// Released under the terms of the GNU Lesser General Public License.
/// </copyright>
#pragma endregion

#include <streams.h>     // Active Movie (includes windows.h)
#include <initguid.h>    // declares DEFINE_GUID to declare an EXTERN_C const.
#include <Dvdmedia.h>
#include <wxdebug.h>

//#include <qedit.h>
#include "SimpleGrabber.h"
#include "DirectShowUtility.h"

#include "openimagelib\il.hpp"

#pragma warning(disable: 4800)

const wchar_t* ScreenshotCommand = L"Screenshot:";
const size_t ScreenshotCommandLength = 11;

const AMOVIESETUP_MEDIATYPE sudPinTypes =
{
    &MEDIATYPE_Video,       // Major type
    &MEDIASUBTYPE_NULL      // Minor type
};

const AMOVIESETUP_PIN psudSimpleGrabberPins[] =
{
	{
		L"In",            // strName
		FALSE,               // bRendered
		FALSE,               // bOutput
		FALSE,               // bZero
		FALSE,               // bMany
		&CLSID_NULL,         // clsConnectsToFilter
		L"",                 // strConnectsToPin
		1,                   // nTypes
		&sudPinTypes,        // lpTypes
	},
	{
		L"Out",           // strName
		FALSE,               // bRendered
		TRUE,                // bOutput
		FALSE,               // bZero
		FALSE,               // bMany
		&CLSID_NULL,         // clsConnectsToFilter
		L"",                 // strConnectsToPin
		1,                   // nTypes
		&sudPinTypes         // lpTypes
	}
};

const AMOVIESETUP_FILTER sudSimpleGrabber =
{
	&CLSID_SimpleGrabber,    // clsID
	L"SimpleGrabber",        // strName
	MERIT_DO_NOT_USE,        // dwMerit MERIT_DO_NOT_USE, 0x00900000
	2,                       // nPins
	psudSimpleGrabberPins    // lpPin
};

// Needed for the CreateInstance mechanism
CFactoryTemplate g_Templates[]=
{
    { L"Simple Grabber"
        , &CLSID_SimpleGrabber
        , CSimpleGrabber::CreateInstance
        , NULL
        , &sudSimpleGrabber }

};

int g_cTemplates = sizeof(g_Templates)/sizeof(g_Templates[0]);


////////////////////////////////////////////////////////////////////////
//
// Exported entry points for registration and unregistration 
// (in this case they only call through to default implementations).
//
////////////////////////////////////////////////////////////////////////

STDAPI DllRegisterServer() 
{
    return AMovieDllRegisterServer2(TRUE);
}

STDAPI DllUnregisterServer() 
{
    return AMovieDllRegisterServer2(FALSE);
}


////////////////////////////////////////////////////////////////////////
//
// DllEntryPoint
//
////////////////////////////////////////////////////////////////////////
extern "C" BOOL WINAPI DllEntryPoint(HINSTANCE, ULONG, LPVOID);

BOOL APIENTRY DllMain(HANDLE hModule, DWORD  dwReason, LPVOID lpReserved)
{
    return DllEntryPoint((HINSTANCE)(hModule), dwReason, lpReserved);
}




////////////////////////////////////////////////////////////////////////
//
// CSimpleGrabberMemAllocator
//
////////////////////////////////////////////////////////////////////////

STDMETHODIMP CSimpleGrabberMemAllocator::GetBuffer(IMediaSample** ppBuffer,
	REFERENCE_TIME* pStartTime, REFERENCE_TIME* pEndTime, DWORD dwFlags)
{
	HRESULT hr = CMemAllocator::GetBuffer(ppBuffer, pStartTime, pEndTime, dwFlags);
	if(FAILED(hr)) return hr;

	if(_mediaTypeChanged && !_upstreamFilterNotified)
	{
		CAutoLock lock(this);

		hr = (*ppBuffer)->SetMediaType(&_mediaType);
		_upstreamFilterNotified = true;
	}
	return hr;
}

HRESULT CSimpleGrabberMemAllocator::GetMediaType(CMediaType& mediaType) const
{
	CAutoLock lock(const_cast<CSimpleGrabberMemAllocator*>(this));

	if(!_mediaTypeChanged || !_upstreamFilterNotified)
		return S_FALSE;

	return mediaType.Set(_mediaType);
}

HRESULT CSimpleGrabberMemAllocator::SetMediaType(const AM_MEDIA_TYPE& mt)
{
	ALLOCATOR_PROPERTIES props;
	HRESULT hr = this->GetProperties(&props);
	if(FAILED(hr)) return hr;

    if(mt.formattype == FORMAT_VideoInfo)
    {
		if(mt.cbFormat < sizeof(VIDEOINFOHEADER) || mt.pbFormat == NULL)
			return E_INVALIDARG;

		VIDEOINFOHEADER* vih = (VIDEOINFOHEADER*)mt.pbFormat;
		props.cbBuffer = vih->bmiHeader.biSizeImage;
	}
	else if(mt.formattype == FORMAT_VideoInfo2)
	{
		if(mt.cbFormat < sizeof(VIDEOINFOHEADER2) || mt.pbFormat == NULL)
			return E_INVALIDARG;

		VIDEOINFOHEADER2* vih = (VIDEOINFOHEADER2*)mt.pbFormat;
		props.cbBuffer = vih->bmiHeader.biSizeImage;
	}
	else
		return E_INVALIDARG;

	ALLOCATOR_PROPERTIES newProps;
	hr = this->SetProperties(&props, &newProps);
	if(FAILED(hr)) return hr;

    CAutoLock lock(this);

	hr = _mediaType.Set(mt);
	if(FAILED(hr)) return hr;

	_mediaTypeChanged = true;
	_upstreamFilterNotified = false;

	return hr;
}




////////////////////////////////////////////////////////////////////////
//
// CSimpleGrabber
//
////////////////////////////////////////////////////////////////////////

//
// Provide the way for COM to create a CSimpleGrabber object
//
CUnknown* WINAPI CSimpleGrabber::CreateInstance(LPUNKNOWN punk, HRESULT *phr) 
{
    ASSERT(phr);
    
    // assuming we don't want to modify the data
    CSimpleGrabber *pNewObject = new CSimpleGrabber(punk, phr);

    if(pNewObject == NULL) {
        if (phr)
            *phr = E_OUTOFMEMORY;
    }

    return pNewObject;   

}


//----------------------------------------------------------------------------
//
//----------------------------------------------------------------------------

CSimpleGrabber::CSimpleGrabber(IUnknown* pOuter, HRESULT* phr)
	: CTransformFilter(TEXT("SimpleGrabber"), (IUnknown*) pOuter, CLSID_SimpleGrabber),
		_commandMessageSink(this), _takeScreenshotState(STATE_WAITING),
		_takeScreenshotRequestTime(0), _sampleBuffer(NULL), _sampleBufferSize(0)
{
    m_pInput = static_cast<CTransformInputPin*>(new CSimpleGrabberInPin(this, phr));
    if(!m_pInput)
    {
        if(phr)
            *phr = E_OUTOFMEMORY;
    }

	m_pOutput = static_cast<CTransformOutputPin*>(new CSimpleGrabberOutPin(this, phr));
    if(!m_pOutput)
    {
        if(phr)
            *phr = E_OUTOFMEMORY;
    }

	_commandMessageSink.Create(NULL, NULL, L"SimpleGrabber_CmdParser");
	ATLASSERT(_commandMessageSink);
}

CSimpleGrabber::~CSimpleGrabber()
{
	if(_commandMessageSink)
	{
		// Make sure we are destroying the window from the right thread.
		_commandMessageSink.SendMessageW(WM_APP_SHUTDOWN);
	}

	if(_sampleBuffer != NULL)
	{
		delete[] _sampleBuffer;
		_sampleBuffer = NULL;
	}
}

CSimpleGrabberInPin* CSimpleGrabber::GetInput()
{
	return static_cast<CSimpleGrabberInPin*>(m_pInput);
}

CSimpleGrabberOutPin* CSimpleGrabber::GetOutput()
{
	return static_cast<CSimpleGrabberOutPin*>(m_pOutput);
}

STDMETHODIMP CSimpleGrabber::NonDelegatingQueryInterface( REFIID riid, void ** ppv) 
{
    CheckPointer(ppv,E_POINTER);

	return CTransformFilter::NonDelegatingQueryInterface(riid, ppv);
}

STDMETHODIMP CSimpleGrabber::JoinFilterGraph(IFilterGraph* pGraph, LPCWSTR pName)
{
	if(pGraph && !m_pGraph)
	{
		HRESULT hr = DirectShowUtility::FilterExists(pGraph, CLSID_SimpleGrabber, this);
		if(hr == S_OK)
			return E_FAIL;
	}

	return CTransformFilter::JoinFilterGraph(pGraph, pName);
}


//----------------------------------------------------------------------------
// This is where you force the sample grabber to connect with one type
// or the other. What you do here is crucial to what type of data your
// app will be dealing with.
//----------------------------------------------------------------------------
HRESULT CSimpleGrabber::CheckInputType(const CMediaType* pmt)
{
	CheckPointer(pmt, E_POINTER);

	if(*pmt->FormatType() != FORMAT_VideoInfo && *pmt->FormatType() != FORMAT_VideoInfo2)
		return E_INVALIDARG;

	if(*pmt->Type() != MEDIATYPE_Video)
		return VFW_E_INVALID_MEDIA_TYPE;

	// Check subtype
	GUID g = *pmt->Subtype();
	if(g == MEDIASUBTYPE_RGB32 ||
		g == MEDIASUBTYPE_ARGB32 || 
		g == MEDIASUBTYPE_YV12 ||
		g == MEDIASUBTYPE_YUY2)
	{
		return NOERROR;
	}
	else
		return VFW_E_INVALID_MEDIA_TYPE;

	return S_OK;
}

HRESULT CSimpleGrabber::QueryAcceptOutput(const AM_MEDIA_TYPE* pmt)
{
    CheckPointer(pmt, E_INVALIDARG);
	if(m_pInput == NULL || m_pInput->IsConnected() == FALSE) return E_INVALIDARG;

	IPin* pin = m_pInput->GetConnected();
	HRESULT hr = pin->QueryAccept(pmt);
	if(FAILED(hr)) return hr;

	CSimpleGrabberMemAllocator* all = GetInput()->GetAllocator();
	if(all == NULL) return E_INVALIDARG;

	hr = all->SetMediaType(*pmt);
	if(FAILED(hr)) return hr;

	return S_OK;
}


HRESULT CSimpleGrabber::AreMediaTypesCompatible(const AM_MEDIA_TYPE* mtIn, const AM_MEDIA_TYPE* mtOut)
{
    CheckPointer(mtIn, E_POINTER);
    CheckPointer(mtOut, E_POINTER);

    // format must be a VIDEOINFOHEADER
    if(mtIn->formattype == FORMAT_VideoInfo &&
		mtOut->formattype == FORMAT_VideoInfo)
    {
		// formats must be big enough 
		if(mtIn->cbFormat < sizeof(VIDEOINFOHEADER) ||
			mtOut->cbFormat < sizeof(VIDEOINFOHEADER))
			return S_FALSE;

		VIDEOINFOHEADER* pInput  = (VIDEOINFOHEADER*)mtIn->pbFormat;
		VIDEOINFOHEADER* pOutput = (VIDEOINFOHEADER*)mtOut->pbFormat;

		if(pInput->bmiHeader.biBitCount != pOutput->bmiHeader.biBitCount ||
			pInput->bmiHeader.biCompression != pOutput->bmiHeader.biCompression)
			return S_FALSE;

		//if(memcmp(&pInput->bmiHeader, &pOutput->bmiHeader, sizeof(BITMAPINFOHEADER)) == 0)
		//	return NOERROR;

		//if(pInput->bmiHeader.biHeight == -pOutput->bmiHeader.biHeight)
		//{
		//	BITMAPINFOHEADER bihOut;
		//	memcpy(&bihOut, &pOutput->bmiHeader, sizeof(BITMAPINFOHEADER));
		//	bihOut.biHeight = -bihOut.biHeight;

		//	if(memcmp(&pInput->bmiHeader, &bihOut, sizeof(BITMAPINFOHEADER)) == 0)
		//		return NOERROR;
		//}
		//return S_FALSE;
    }
/*	else if(mtIn->formattype == FORMAT_VideoInfo &&
		mtOut->formattype == FORMAT_VideoInfo2)
    {
		// formats must be big enough 
		if(mtIn->cbFormat < sizeof(VIDEOINFOHEADER) ||
			mtOut->cbFormat < sizeof(VIDEOINFOHEADER2))
			return S_FALSE;

		VIDEOINFOHEADER* pInput  = (VIDEOINFOHEADER*)mtIn->pbFormat;
		VIDEOINFOHEADER2* pOutput = (VIDEOINFOHEADER2*)mtOut->pbFormat;

		if(pInput->bmiHeader.biBitCount != pOutput->bmiHeader.biBitCount ||
			pInput->bmiHeader.biCompression != pOutput->bmiHeader.biCompression)
			return S_FALSE;

		//if(memcmp(&pInput->bmiHeader, &pOutput->bmiHeader, sizeof(BITMAPINFOHEADER)) == 0)
		//	return NOERROR;

		//if(pInput->bmiHeader.biHeight == -pOutput->bmiHeader.biHeight)
		//{
		//	BITMAPINFOHEADER bihOut;
		//	memcpy(&bihOut, &pOutput->bmiHeader, sizeof(BITMAPINFOHEADER));
		//	bihOut.biHeight = -bihOut.biHeight;

		//	if(memcmp(&pInput->bmiHeader, &bihOut, sizeof(BITMAPINFOHEADER)) == 0)
		//		return NOERROR;
		//}
		//return S_FALSE;
	}
*/
	else if(mtIn->formattype == FORMAT_VideoInfo2 &&
		mtOut->formattype == FORMAT_VideoInfo2)
    {
		// formats must be big enough 
		if(mtIn->cbFormat < sizeof(VIDEOINFOHEADER2) ||
			mtOut->cbFormat < sizeof(VIDEOINFOHEADER2))
			return S_FALSE;

		VIDEOINFOHEADER2* pInput = (VIDEOINFOHEADER2*)mtIn->pbFormat;
		VIDEOINFOHEADER2* pOutput = (VIDEOINFOHEADER2*)mtOut->pbFormat;

		if(pInput->bmiHeader.biBitCount != pOutput->bmiHeader.biBitCount ||
			pInput->bmiHeader.biCompression != pOutput->bmiHeader.biCompression)
			return S_FALSE;

		//if(memcmp(&pInput->bmiHeader, &pOutput->bmiHeader, sizeof(BITMAPINFOHEADER)) != 0)
		//	return S_FALSE;

		//if(pInput->dwInterlaceFlags != pOutput->dwInterlaceFlags ||
		//	pInput->dwControlFlags != pOutput->dwControlFlags)
		//	return S_FALSE;
	}
	else
		return S_FALSE;

	return NOERROR;
}

//
// CheckTransform
//
// To be able to transform the formats must be identical
//
HRESULT CSimpleGrabber::CheckTransform(const CMediaType* mtIn, const CMediaType* mtOut)
{
    CheckPointer(mtIn, VFW_E_TYPE_NOT_ACCEPTED);
    CheckPointer(mtOut, VFW_E_TYPE_NOT_ACCEPTED);

	HRESULT hr = CheckInputType(mtIn);
    if(FAILED(hr)) return VFW_E_TYPE_NOT_ACCEPTED;

	hr = CSimpleGrabber::AreMediaTypesCompatible(mtIn, mtOut);
	if(hr == S_OK) return hr;

	if(m_pInput == NULL || m_pInput->IsConnected() == FALSE)
		return VFW_E_TYPE_NOT_ACCEPTED;

	hr = GetInput()->CanReconnect(*mtOut, NULL);
	if(hr == S_OK) return hr;

	return VFW_E_TYPE_NOT_ACCEPTED;
}

//
// DecideBufferSize
//
// Tell the output pin's allocator what size buffers we
// require. Can only do this when the input is connected
//
HRESULT CSimpleGrabber::DecideBufferSize(IMemAllocator* pAlloc, ALLOCATOR_PROPERTIES* pProperties)
{
    CheckPointer(pAlloc,E_POINTER);
    CheckPointer(pProperties,E_POINTER);

    // Is the input pin connected
	if(m_pInput->IsConnected() == FALSE)
		return E_UNEXPECTED;

    HRESULT hr = NOERROR;
    pProperties->cBuffers = 1;
    pProperties->cbBuffer = m_pInput->CurrentMediaType().GetSampleSize();

    ASSERT(pProperties->cbBuffer);

    // If we don't have fixed sized samples we must guess some size
    if(!m_pInput->CurrentMediaType().bFixedSizeSamples)
    {
        if(pProperties->cbBuffer < 100000)
        {
            // nothing more than a guess!!
            pProperties->cbBuffer = 100000;
        }
    }

    // Ask the allocator to reserve us some sample memory, NOTE the function
    // can succeed (that is return NOERROR) but still not have allocated the
    // memory that we requested, so we must check we got whatever we wanted

    ALLOCATOR_PROPERTIES Actual;
    hr = pAlloc->SetProperties(pProperties,&Actual);
    if(FAILED(hr))
    {
        return hr;
    }

    ASSERT(Actual.cBuffers == 1);

    if(pProperties->cBuffers > Actual.cBuffers ||
        pProperties->cbBuffer > Actual.cbBuffer)
    {
        return E_FAIL;
    }

    return NOERROR;

}

HRESULT CSimpleGrabber::GetMediaType(int iPosition, CMediaType* pMediaType)
{
	HRESULT hr = GetInput()->GetUpFilterMediaType(iPosition, pMediaType);
    return hr;
}

HRESULT CSimpleGrabber::CompleteConnect(PIN_DIRECTION direction, IPin* pReceivePin)
{
	if(direction == PINDIR_INPUT)
		return S_OK;

	if(pReceivePin == NULL) return E_POINTER;
	if(m_pGraph == NULL) return E_UNEXPECTED;
	if(m_pInput->IsConnected() == FALSE) return E_UNEXPECTED;

	HRESULT hr = CSimpleGrabber::AreMediaTypesCompatible(&m_pInput->CurrentMediaType(), &m_pOutput->CurrentMediaType());
	if(hr == S_OK) return S_OK;

	CMediaType upFilterMediaType;
	hr = GetInput()->CanReconnect(m_pOutput->CurrentMediaType(), &upFilterMediaType);
	if(FAILED(hr)) return hr;
	if(hr == S_FALSE) return VFW_E_TYPE_NOT_ACCEPTED;

	CComPtr<IFilterGraph2> filterGraph2;
	hr = m_pGraph->QueryInterface(&filterGraph2);
	if(FAILED(hr)) return hr;
	
	CComPtr<IPin> upFilterOutputPin;
	hr = m_pInput->ConnectedTo(&upFilterOutputPin);
	if(FAILED(hr)) return hr;
	
	hr = filterGraph2->ReconnectEx(upFilterOutputPin, &upFilterMediaType);

	return hr;
}


//
// Makes destination an identical copy of source
//
HRESULT CSimpleGrabber::Copy(IMediaSample *pSource, IMediaSample *pDest) const
{
    CheckPointer(pSource,E_POINTER);
    CheckPointer(pDest,E_POINTER);

    // Copy the sample data
    BYTE *pSourceBuffer, *pDestBuffer;
    long lSourceSize = pSource->GetActualDataLength();

#ifdef DEBUG
    long lDestSize = pDest->GetSize();
    ASSERT(lDestSize >= lSourceSize);
#endif

    pSource->GetPointer(&pSourceBuffer);
    pDest->GetPointer(&pDestBuffer);

    AM_MEDIA_TYPE* srcMediaType;
    pSource->GetMediaType(&srcMediaType);

    CopyMemory((PVOID) pDestBuffer,(PVOID) pSourceBuffer, lSourceSize);

    // Copy the media type

    HRESULT hr = pDest->SetMediaType(srcMediaType);
	if(FAILED(hr)) return hr;
    
	DeleteMediaType(srcMediaType);

    // Copy the sample times

    REFERENCE_TIME TimeStart, TimeEnd;
    if(NOERROR == pSource->GetTime(&TimeStart, &TimeEnd))
    {
        pDest->SetTime(&TimeStart, &TimeEnd);
    }

    LONGLONG MediaStart, MediaEnd;
    if(pSource->GetMediaTime(&MediaStart,&MediaEnd) == NOERROR)
    {
        pDest->SetMediaTime(&MediaStart,&MediaEnd);
    }

    // Copy the Sync point property

    hr = pSource->IsSyncPoint();
    if(hr == S_OK)
    {
        pDest->SetSyncPoint(TRUE);
    }
    else if(hr == S_FALSE)
    {
        pDest->SetSyncPoint(FALSE);
    }
    else
    {  // an unexpected error has occured...
        return E_UNEXPECTED;
    }


    // Copy the preroll property

    hr = pSource->IsPreroll();
    if(hr == S_OK)
    {
        pDest->SetPreroll(TRUE);
    }
    else if(hr == S_FALSE)
    {
        pDest->SetPreroll(FALSE);
    }
    else
    {  // an unexpected error has occured...
        return E_UNEXPECTED;
    }

    // Copy the discontinuity property

    hr = pSource->IsDiscontinuity();

    if(hr == S_OK)
    {
        pDest->SetDiscontinuity(TRUE);
    }
    else if(hr == S_FALSE)
    {
        pDest->SetDiscontinuity(FALSE);
    }
    else
    {  // an unexpected error has occured...
        return E_UNEXPECTED;
    }

    // Copy the actual data length

    long lDataLength = pSource->GetActualDataLength();
    pDest->SetActualDataLength(lDataLength);

    return NOERROR;

} // Copy



//----------------------------------------------------------------------------
// Transform
//----------------------------------------------------------------------------

HRESULT CSimpleGrabber::Transform(IMediaSample* pIn, IMediaSample* pOut)
{
    CheckPointer(pIn, E_POINTER);
    CAutoLock lock(&m_Lock);

	HRESULT hr = S_OK;
    if(pIn != NULL)
    {
        //REFERENCE_TIME StartTime, StopTime;
        //pms->GetTime( &StartTime, &StopTime);

        //StartTime += m_pInput->CurrentStartTime();
        //StopTime  += m_pInput->CurrentStartTime();
		
		if(_takeScreenshotState == STATE_PREPARING)
		{
			if(::GetTickCount() - _takeScreenshotRequestTime < 1000)
			{
				HRESULT hr2 = PrepareScreenshot(pIn);
				if(SUCCEEDED(hr2))
				{
					_takeScreenshotState = STATE_SAVING;
					_commandMessageSink.PostMessageW(WM_APP_SAVESCREENSHOT);
				}
			}
			else
				_takeScreenshotState = STATE_WAITING;
		}
    }

	hr = Copy(pIn, pOut);

    return hr;
}

LRESULT CSimpleGrabber::OnCommandWindowSetText(UINT uMsg, WPARAM wParam, LPARAM lParam, BOOL& bHandled)
{
	const wchar_t* commandText = (LPCTSTR)lParam;

	//HRESULT hr = S_OK;
	if(_wcsnicmp(commandText, ScreenshotCommand, ScreenshotCommandLength) == 0)
	{
		CAutoLock lock(&m_Lock);

		if(_takeScreenshotState == STATE_WAITING ||
			_takeScreenshotState == STATE_PREPARING)
		{
			_takeScreenshotRequestTime = ::GetTickCount();
			_screenshotFilePath.clear();
			_screenshotFilePath.append(commandText + ScreenshotCommandLength);
			_takeScreenshotState = STATE_PREPARING;
		}
	}

	return FALSE;
}

HRESULT CSimpleGrabber::PrepareScreenshot(IMediaSample* mediaSample)
{
	if(mediaSample == NULL)
		return E_INVALIDARG;

	HRESULT hr = S_OK;

	if(m_pInput == NULL)
		return E_UNEXPECTED;

	CMediaType mediaType;
	CSimpleGrabberMemAllocator* allocator = GetInput()->GetAllocator();
	if(allocator == NULL) return E_UNEXPECTED;

	hr = allocator->GetMediaType(mediaType);
	if(hr == S_FALSE)
		hr = m_pInput->ConnectionMediaType(&mediaType);

	if(FAILED(hr)) return hr;

	if(*mediaType.Type() != MEDIATYPE_Video)
		return E_INVALIDARG;

	long bufferSize = mediaSample->GetActualDataLength();
	if(bufferSize <= 0)
		return E_INVALIDARG;

	BYTE* buffer = NULL;
	hr = mediaSample->GetPointer(&buffer);
	if(FAILED(hr)) return hr;

	if(*mediaType.FormatType() != FORMAT_VideoInfo &&
		*mediaType.FormatType() != FORMAT_VideoInfo2)
		return HRESULT_FROM_WIN32(ERROR_NOT_SUPPORTED);

	// Copy sample media type
	_sampleMediaType = mediaType;

	// Realocate buffer if necessary
	if(_sampleBuffer == NULL || _sampleBufferSize < (size_t)bufferSize)
	{
		if(_sampleBuffer != NULL)
		{
			delete[] _sampleBuffer;
			_sampleBuffer = NULL;
			_sampleBufferSize = 0;
		}

		_sampleBuffer = new BYTE[bufferSize];
		if(_sampleBuffer == NULL) return E_OUTOFMEMORY;

		_sampleBufferSize = bufferSize;
	}

	// Copy sample data
	memcpy(_sampleBuffer, buffer, bufferSize);

	return hr;
}

LRESULT CSimpleGrabber::OnSaveScreenshot(UINT uMsg, WPARAM wParam, LPARAM lParam, BOOL& bHandled)
{
	SaveScreenshot();
	_takeScreenshotState = STATE_WAITING;
	return FALSE;
}

LRESULT CSimpleGrabber::OnSaveShutDown(UINT uMsg, WPARAM wParam, LPARAM lParam, BOOL& bHandled)
{
	_commandMessageSink.DestroyWindow();
	return FALSE;
}


namespace il  = olib::openimagelib::il;
namespace olib { namespace openimagelib { namespace il {

typedef image< unsigned char, l8 >				l8_image_type;
typedef image< unsigned char, l8a8 >			l8a8_image_type;
typedef image< unsigned char, r8g8b8 >			r8g8b8_image_type;
typedef image< unsigned char, b8g8r8 >			b8g8r8_image_type;
typedef image< unsigned char, b8g8r8a8 >		b8g8r8a8_image_type;
typedef image< unsigned char, r8g8b8a8 >		r8g8b8a8_image_type;
typedef image< unsigned char, a8r8g8b8 >		a8r8g8b8_image_type;
typedef image< unsigned char, a8b8g8r8 >		a8b8g8r8_image_type;
typedef image< unsigned char, r16g16b16a16f >	r16g16b16a16f_image_type;
typedef image< unsigned char, yuv444 >			yuv444_image_type;
typedef image< unsigned char, yuv422 >			yuv422_image_type;
typedef image< unsigned char, yuv422p >			yuv422p_image_type;
typedef image< unsigned char, yuv420p >			yuv420p_image_type;
typedef image< unsigned char, yuv411 >			yuv411_image_type;
typedef image< unsigned char, yuv411p >			yuv411p_image_type;

}}}

HRESULT CSimpleGrabber::SaveScreenshot()
{
	HRESULT hr = S_OK;
	RECT sourceRect = {0};
    int width  = 0;
    int height = 0;
	int heightAbs = 0;
	WORD colors = 0;

	//assert(0);

	if(*_sampleMediaType.FormatType() == FORMAT_VideoInfo)
	{
		VIDEOINFOHEADER* vih = (VIDEOINFOHEADER*)_sampleMediaType.Format();
		sourceRect = vih->rcSource;
		width  = vih->bmiHeader.biWidth;
		height = vih->bmiHeader.biHeight;
		heightAbs = (height < 0 ? -height : height);
		colors = vih->bmiHeader.biBitCount;
	}
	else if(*_sampleMediaType.FormatType() == FORMAT_VideoInfo2)
	{
		VIDEOINFOHEADER2* vih = (VIDEOINFOHEADER2*)_sampleMediaType.Format();
		sourceRect = vih->rcSource;
		width  = vih->bmiHeader.biWidth;
		height = vih->bmiHeader.biHeight;
		heightAbs = (height < 0 ? -height : height);
		colors = vih->bmiHeader.biBitCount;
	}
	else
		return E_UNEXPECTED;

	il::image_type_ptr resultImage;
	const BYTE* resultImageBuffer = NULL;
	size_t resultImageBufferSize = 0;
	il::image_type_ptr srcImage;
	// Because of a dynamic format change the actual image size might not match to the size of a bitmap.
	int resultImageWidth = sourceRect.right - sourceRect.left;
	int resultImageHeight = sourceRect.bottom - sourceRect.top;

	if(*_sampleMediaType.Subtype() == MEDIASUBTYPE_RGB24)
	{
		resultImageBuffer = _sampleBuffer;
		resultImageBufferSize = _sampleBufferSize;
	}
	else
	{
		bool flip = false;
		if(*_sampleMediaType.Subtype() == MEDIASUBTYPE_RGB32 ||
			*_sampleMediaType.Subtype() == MEDIASUBTYPE_ARGB32)
		{
			srcImage = il::image_type_ptr(new il::image_type(il::b8g8r8a8_image_type(width, heightAbs, 1)));
			flip = height < 0;
		}
		else if(*_sampleMediaType.Subtype() == MEDIASUBTYPE_YV12)
		{
			srcImage = il::image_type_ptr(new il::image_type(il::yuv420p_image_type(width, heightAbs, 1)));
			flip = true;
		}
		else if(*_sampleMediaType.Subtype() == MEDIASUBTYPE_YUY2)
		{
			srcImage = il::image_type_ptr(new il::image_type(il::yuv422_image_type(width, heightAbs, 1)));
			flip = true;
		}
		else
			return E_NOTIMPL;

		srcImage->crop(sourceRect.left, sourceRect.top, resultImageWidth, resultImageHeight);

		srcImage->get_structure()->destroy();
		srcImage->get_structure()->set_data(_sampleBuffer, _sampleBufferSize);
		resultImage = il::convert(srcImage, L"b8g8r8");
		srcImage->get_structure()->set_data(NULL, 0);

		assert(resultImage != NULL);
		if(resultImage != NULL)
		{
			if(flip)
				resultImage = il::image_type_ptr(resultImage->clone(il::flipped));

			resultImageBuffer = resultImage->data(0);
			resultImageBufferSize = resultImage->size();
		}
		else
			hr = E_FAIL;
	}

	if(FAILED(hr)) return hr;

	hr = SaveBitmapFile(_screenshotFilePath.c_str(), resultImageBuffer, resultImageBufferSize,
		resultImageWidth, resultImageHeight, 24);

	return hr;
}

HRESULT CSimpleGrabber::SaveBitmapFile(LPCWSTR filePath,
	const BYTE* bitmapBits, size_t bitmapBitsSize, int width, int height, WORD colorMode)
{
	HRESULT hr = S_OK;

	BITMAPFILEHEADER bmpFileHeader = {0};
	bmpFileHeader.bfType = 'MB';
	bmpFileHeader.bfSize = sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER) + (DWORD)bitmapBitsSize;
	bmpFileHeader.bfOffBits = sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER);

	BITMAPINFOHEADER bmpInfoHeader = {0};
	bmpInfoHeader.biSize = sizeof(bmpInfoHeader);
	bmpInfoHeader.biWidth = width;
	bmpInfoHeader.biHeight = height;
	bmpInfoHeader.biPlanes = 1;
	bmpInfoHeader.biBitCount = colorMode;

	// Create the .BMP file.
	HANDLE hf = ::CreateFile(filePath, GENERIC_READ | GENERIC_WRITE, 0, NULL,
		CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
	if(hf == INVALID_HANDLE_VALUE)
		return HRESULT_FROM_WIN32(::GetLastError());

	do
	{
		DWORD written = 0;
		if(!::WriteFile(hf, &bmpFileHeader, sizeof(bmpFileHeader), &written, NULL))
		{
			hr = HRESULT_FROM_WIN32(::GetLastError());
			break;
		}

		if(!::WriteFile(hf, &bmpInfoHeader, sizeof(bmpInfoHeader), &written, NULL))
		{
			hr = HRESULT_FROM_WIN32(::GetLastError());
			break;
		}

		if(!::WriteFile(hf, bitmapBits, (DWORD)bitmapBitsSize, &written, NULL))
			hr = HRESULT_FROM_WIN32(::GetLastError());
	}
	while(!this);

	// Close the .BMP file. 
	if(!CloseHandle(hf)) 
		hr = HRESULT_FROM_WIN32(::GetLastError());

	return hr;
}

HRESULT CSimpleGrabber::SaveBitmapFile(LPCWSTR filePath, const BITMAPINFO* bitmapInfo)
{
	PBITMAPINFOHEADER pbih = (PBITMAPINFOHEADER) bitmapInfo;
	LPBYTE lpBits = (LPBYTE)((LPSTR)bitmapInfo + (WORD)(bitmapInfo->bmiHeader.biSize));

	BITMAPFILEHEADER hdr;       // bitmap file-header
	hdr.bfType = 0x4d42;        // 0x42 = "B" 0x4d = "M"
	// Compute the size of the entire file.
	hdr.bfSize = (DWORD) (sizeof(BITMAPFILEHEADER) + pbih->biSize + 
		pbih->biClrUsed * sizeof(RGBQUAD) + pbih->biSizeImage);
	hdr.bfReserved1 = 0;
	hdr.bfReserved2 = 0;

	// Compute the offset to the array of color indices. 
	hdr.bfOffBits = (DWORD) sizeof(BITMAPFILEHEADER) + pbih->biSize + 
		pbih->biClrUsed * sizeof (RGBQUAD); 


	HRESULT hr = S_OK;

	// Create the .BMP file.
	HANDLE hf = ::CreateFile(filePath, GENERIC_READ | GENERIC_WRITE, 0, NULL,
		CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
    if(hf == INVALID_HANDLE_VALUE)
		return HRESULT_FROM_WIN32(::GetLastError());

	// Copy the BITMAPFILEHEADER into the .BMP file. 
	DWORD dwTmp;
	if(::WriteFile(hf, &hdr, sizeof(BITMAPFILEHEADER), &dwTmp,  NULL))
	{
	    // Copy the BITMAPINFOHEADER and RGBQUAD array into the file. 
		if(::WriteFile(hf, pbih, sizeof(BITMAPINFOHEADER) + pbih->biClrUsed * sizeof (RGBQUAD),
			(LPDWORD) &dwTmp, NULL))
		{
			// Copy the array of color indices into the .BMP file. 
			DWORD cb = pbih->biSizeImage; 
			if(!::WriteFile(hf, lpBits, cb, &dwTmp, NULL))
				hr = HRESULT_FROM_WIN32(::GetLastError());
		}
		else
			hr = HRESULT_FROM_WIN32(::GetLastError());
	}
	else
		hr = HRESULT_FROM_WIN32(::GetLastError());


	// Close the .BMP file. 
	if(!CloseHandle(hf)) 
		hr = HRESULT_FROM_WIN32(::GetLastError());

	return hr;
}




////////////////////////////////////////////////////////////////////////
//
// CSimpleGrabberOutPin
//
////////////////////////////////////////////////////////////////////////


CSimpleGrabberInPin::~CSimpleGrabberInPin()
{
	ClearMediaTypes();
}

void CSimpleGrabberInPin::ClearMediaTypes()
{
	CAutoLock lock(m_pLock);

	for(size_t i = 0; i < _upFilterMediaTypes.size(); i++)
	{
		AM_MEDIA_TYPE* mediaType = _upFilterMediaTypes[i];
		if(mediaType != NULL)
			DeleteMediaType(mediaType);
	}

	_upFilterMediaTypes.clear();
}

HRESULT CSimpleGrabberInPin::LoadMediaTypes(IPin* pin)
{
	CAutoLock lock(m_pLock);

	CComPtr<IEnumMediaTypes> mediaTypes;
	HRESULT hr = pin->EnumMediaTypes(&mediaTypes);
	if(FAILED(hr)) return hr;

	ULONG n;
	AM_MEDIA_TYPE* mediaTypeArray = NULL;
	while((hr = mediaTypes->Next(1, &mediaTypeArray, &n)) == S_OK)
	{
		if(mediaTypeArray == NULL) continue;
		_upFilterMediaTypes.push_back(mediaTypeArray);
	}
	if(hr == S_FALSE) return S_OK;
	return hr;
}

HRESULT CSimpleGrabberInPin::CanReconnect(const CMediaType& outputMediaType, CMediaType* upFilterMediaType)
{
    if(!IsStopped() && !m_bCanReconnectWhenActive)
        return S_FALSE;

	CAutoLock lock(m_pLock);

	for(size_t i = 0; i < _upFilterMediaTypes.size(); i++)
	{
		AM_MEDIA_TYPE* nextMediaType = _upFilterMediaTypes[i];
		HRESULT hr = CSimpleGrabber::AreMediaTypesCompatible(nextMediaType, &outputMediaType);
		if(hr == S_OK)
		{
			if(upFilterMediaType != NULL)
				(*upFilterMediaType).Set(*nextMediaType);
			return S_OK;
		}
	}

	return S_FALSE;
}

HRESULT CSimpleGrabberInPin::GetUpFilterMediaType(int iPosition, CMediaType* pMediaType)
{
	CAutoLock lock(m_pLock);

	if(IsConnected() == FALSE) return E_UNEXPECTED;
    if(iPosition < 0) return E_INVALIDARG;
	if(iPosition >= (int)_upFilterMediaTypes.size()) return E_INVALIDARG;

	CheckPointer(pMediaType, E_POINTER);
    ValidateReadWritePtr(pMediaType, sizeof(AM_MEDIA_TYPE));

	*pMediaType = CMediaType(*_upFilterMediaTypes[iPosition]);

	return S_OK;	
}

HRESULT CSimpleGrabberInPin::GetMediaType(int iPosition, CMediaType* pMediaType)
{
    if(iPosition < 0) return E_INVALIDARG;
    if(iPosition > 0) return VFW_S_NO_MORE_ITEMS;

	CheckPointer(pMediaType, E_POINTER);
    ValidateReadWritePtr(pMediaType, sizeof(AM_MEDIA_TYPE));
	
	*pMediaType = CMediaType();
	pMediaType->SetType(&MEDIATYPE_Video);
	
    return S_OK;
}

STDMETHODIMP CSimpleGrabberInPin::GetAllocator(IMemAllocator** ppAllocator)
{
	CheckPointer(ppAllocator,E_POINTER);
	ValidateReadWritePtr(ppAllocator,sizeof(IMemAllocator *));
	CAutoLock lock(m_pLock);

	if(m_pAllocator == NULL)
	{
		HRESULT hr = S_OK;
		CUnknown* unk = new CSimpleGrabberMemAllocator(&hr);
		if(FAILED(hr)) return hr;
		if(unk == NULL) return E_OUTOFMEMORY;

		hr = unk->NonDelegatingQueryInterface(__uuidof(IMemAllocator), reinterpret_cast<void**>(&m_pAllocator));
		if(FAILED(hr)) return hr;
	}

	ASSERT(m_pAllocator != NULL);
	*ppAllocator = m_pAllocator;
	m_pAllocator->AddRef();
	return NOERROR;
}

STDMETHODIMP CSimpleGrabberInPin::NotifyAllocator(IMemAllocator* pAllocator, BOOL bReadOnly)
{
	CheckPointer(pAllocator,E_POINTER);
	ValidateReadPtr(pAllocator,sizeof(IMemAllocator));
	CAutoLock lock(m_pLock);

	if(m_pAllocator != pAllocator)
		return E_INVALIDARG;

	m_bReadOnly = (BYTE)bReadOnly;
	return NOERROR;
}

STDMETHODIMP CSimpleGrabberInPin::ReceiveConnection(IPin* pConnector, const AM_MEDIA_TYPE* pmt)
{
	if(pConnector == NULL) return E_POINTER;
	if(pmt == NULL) return E_POINTER;
    if(m_Connected) return VFW_E_ALREADY_CONNECTED;

	CMediaType mt(*pmt);
	DisplayTypeInfo(pConnector, &mt);

	// Load media types supported by the up-stream filter
	ClearMediaTypes();
	HRESULT hr = LoadMediaTypes(pConnector);
	if(FAILED(hr)) return hr;
	
	hr = CTransformInputPin::ReceiveConnection(pConnector, pmt);
	return hr;
}


////////////////////////////////////////////////////////////////////////
//
// CSimpleGrabberOutPin
//
////////////////////////////////////////////////////////////////////////

STDMETHODIMP CSimpleGrabberOutPin::QueryAccept(const AM_MEDIA_TYPE* pmt)
{
	HRESULT hr = static_cast<CSimpleGrabber*>(m_pFilter)->QueryAcceptOutput(pmt);
	if(FAILED(hr)) return S_FALSE;
	
	return CTransformOutputPin::QueryAccept(pmt);
}
